package com.tsy.server.component.websocket.config;

import com.tsy.server.component.security.SecurityUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.security.Principal;
import java.util.List;
import java.util.Map;

/**
 * WebSocket认证处理器
 * 用于在WebSocket握手阶段设置Principal，确保连接的安全性
 * 
 * @author sean.zhou
 * @date 2021/11/16
 * @version 0.1
 */
@Slf4j
@Component
public class AuthPrincipalHandler extends DefaultHandshakeHandler {

    /**
     * 确定WebSocket连接的Principal
     * Principal的名称格式为：{workspaceId}/{userType}/{userId}
     * 
     * @param request HTTP请求对象，用于获取认证信息
     * @param wsHandler WebSocket处理器
     * @param attributes 握手属性，可用于存储额外信息
     * @return Principal对象，包含用户身份信息
     */
    @Override
    protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        if (request instanceof ServletServerHttpRequest) {
            try {
                String workspaceId = SecurityUtils.getWorkspaceId();
                String username = SecurityUtils.getCurrentUsername();
                
                // 验证必要的用户信息是否存在
                if (workspaceId == null || username == null) {
                    log.warn("WebSocket连接缺少必要的认证信息: workspaceId={}, userId={}", workspaceId, username);
                    return new StompPrincipal(null);
                }

                // todo userType 字段处理
                // 创建Principal，格式为：workspaceId/userType/userId
                String principalName = workspaceId + "/" + 1 + "/" + username;
                log.debug("WebSocket连接已认证: {}", principalName);
                return new StompPrincipal(principalName);
            } catch (Exception e) {
                log.error("WebSocket认证过程中发生错误", e);
                return new StompPrincipal(null);
            }
        }
        log.warn("非ServletServerHttpRequest类型的请求");
        return new StompPrincipal(null);
    }

    /**
     * 重写选择子协议的方法
     * 该方法用于在WebSocket握手过程中选择使用的子协议
     *
     * @param requestedProtocols 客户端请求的子协议列表
     * @param webSocketHandler WebSocket处理器
     * @return 选择的子协议，如果返回null则表示不使用子协议
     */
    @Override
    protected String selectProtocol(List<String> requestedProtocols, WebSocketHandler webSocketHandler) {
        if (requestedProtocols != null && !requestedProtocols.isEmpty()) {
            // 获取客户端请求的第一个协议
            String protocol = requestedProtocols.get(0);
            if (StringUtils.hasText(protocol)) {
                log.debug("选择WebSocket子协议: {}", protocol);
                // 直接返回客户端请求的第一个协议
                return protocol;
            }
        }
        // 默认返回null，表示不使用子协议
        return null;
    }
}